
#include "stun_protocol_parser.h"

#include "./../string_util/string_util.h"
#include "stun_protocol_attr.h"
#include "stun_protocol_msg.h"

#include <cstdio>
#include <cstring>
#include <memory>

StunProtocolParser::StunProtocolParser()
{
    m_stun_msg = nullptr;
}

void StunProtocolParser::ResetStunMsg()
{
    m_stun_msg = nullptr;
}

std::shared_ptr<StunMsg> StunProtocolParser::StunMsgPtr()
{
    return m_stun_msg;
}

int StunProtocolParser::ComposeStunMsg(uint8_t * data, int data_length) const
{
    int len = ComposeStunMsgHeader(data);
    len    += ComposeStunMsgAttr(data + len);
    return len;
}

int StunProtocolParser::ComposeStunMsgHeader(uint8_t * data)
{
    return 0;
}

int StunProtocolParser::ComposeStunMsgAttr(uint8_t * data) const
{
    int    index = 0;
    auto & attr  = m_stun_msg->Attr();
    for (auto & item : attr)
    {
        index += item->ComposeAttrToBuffer(data + index, 0);
    }
    return index;
}

int StunProtocolParser::ParserStunMsg(const uint8_t * data, int data_length, int & use_length)
{
    if (data_length < 4)
    {
        return NEED_DATA;
    }
    int header_length = 0;
    int attr_length   = 0;
    if (ParserStunMsgHeader(data, data_length, header_length) == FINISH)
    {
        use_length += header_length;
        m_stun_msg->DumpMsgHeader();
        if (ParserStunMsgAttr(data + header_length, data_length - header_length, attr_length) == FINISH)
        {
            use_length += attr_length;
            return FINISH;
        }
        return BAD;
    }
    return BAD;
}

int StunProtocolParser::ParserStunMsgHeader(const uint8_t * data, int data_length, int & use_length)
{
    auto        msg                                                 = std::make_shared<StunMsg>();
    uint16_t    message_type                                        = 0;
    uint16_t    message_len                                         = 0;
    uint32_t    message_cookie                                      = 0;
    uint8_t     message_transaction_id[StunMsg::TRANSACTION_LENGTH] = {};
    std::string transaction;
    int         index = 0;

    memcpy(&message_type, data + index, sizeof(message_type));
    msg->MsgType(StringUtil::BeToH16(message_type));
    index += sizeof(message_type);
    memcpy(&message_len, data + index, sizeof(message_len));
    msg->MsgLength(StringUtil::BeToH16(message_len));
    index += sizeof(message_len);
    if (data_length < index + msg->MsgLength())
    {
        return NEED_DATA;
    }
    memcpy(&message_cookie, data + index, sizeof(message_cookie));
    msg->MsgMagicCookie(StringUtil::BeToH32(message_cookie));
    index += sizeof(message_cookie);
    memcpy(message_transaction_id, data + index, sizeof(message_transaction_id));
    index      += sizeof(message_transaction_id);
    use_length += index;
    StringUtil::DataToHexStr(message_transaction_id, sizeof(message_transaction_id), transaction);
    msg->TransactionId(transaction);
    msg->TransactionId(message_transaction_id);
    m_stun_msg = msg;
    return FINISH;
}

int StunProtocolParser::ParserStunMsgAttr(const uint8_t * data, int data_length, int & use_length) const
{
    auto msg = m_stun_msg;
    if (msg->MsgLength() <= data_length)
    {
        if (msg->MsgLength() == 0)
        {
            return FINISH;
        }
        int index = 0;
        while (true)
        {
            uint16_t attr_type   = 0;
            uint16_t attr_length = 0;
            memcpy(&attr_type, data + index, sizeof(attr_type));
            index    += sizeof(attr_type);
            attr_type = StringUtil::BeToH16(attr_type);
            memcpy(&attr_length, data + index, sizeof(attr_length));
            index      += sizeof(attr_length);
            attr_length = StringUtil::BeToH16(attr_length);
            auto attr   = StunAttr::CreateStunAttr(attr_type);
            if (attr != nullptr)
            {
                int length = 0;
                attr->AttrLength(attr_length);
                attr->MagicCookie(msg->MsgCookie());
                attr->TransactionID(msg->TransactionId(), StunMsg::TRANSACTION_LENGTH);
                int result = attr->ParseBufferToAttr(data + index, data_length - index, length);
                if (result == StunAttr::EmParseSuccess)
                {
                    index += attr->AttrPadding();
                    msg->AddAttr(attr);
                    attr->DumpAttr();
                }
                else
                {
                    break;
                }
            }
            index += attr_length;
            if (index >= msg->MsgLength())
            {
                return FINISH;
            }
        }
    }
    return NEED_DATA;
}
